#!/usr/bin/env python
# -*- coding: UTF-8 -*-
'''
@Project ：MicInter_IPCamera 
@File    ：初始代码.py
@IDE     ：PyCharm 
@Author  ：董文捷
@Date    ：2025/11/5 15:53 
'''
import cv2
import numpy as np
import os
import time
import socket
import threading
import queue
from ultralytics import YOLO
from tqdm import tqdm
from scipy.signal import find_peaks
import matplotlib.pyplot as plt


class MotorController:
    def __init__(self):
        self.running = False
        self.command_queue = queue.Queue()
        self.server_thread = None
        self.sender_thread = None
        self.clients = []
        self.lock = threading.Lock()

    def start_server(self, host="0.0.0.0", port=8000):
        self.server_thread = threading.Thread(target=self._run_server, args=(host, port))
        self.server_thread.daemon = True
        self.server_thread.start()

        self.sender_thread = threading.Thread(target=self._send_to_all_clients)
        self.sender_thread.daemon = True
        self.sender_thread.start()

    def _run_server(self, host, port):
        server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        server_socket.bind((host, port))
        server_socket.listen(5)
        print(f"电机控制服务器正在监听, IP地址为:{socket.gethostbyname(socket.gethostname())}，端口号为：{port}")

        while True:
            client_socket, client_address = server_socket.accept()
            client_handler = threading.Thread(target=self._handle_client, args=(client_socket, client_address))
            client_handler.daemon = True
            client_handler.start()

    def _handle_client(self, client_socket, client_address):
        print(f"电机客户端 {client_address} 已连接")
        with self.lock:
            self.clients.append(client_socket)
        try:
            while True:
                data = client_socket.recv(1024).decode('utf-8')
                if not data:
                    print(f"电机客户端 {client_address} 断开连接")
                    break
                print(f"从电机客户端 {client_address} 收到数据：{data}")
        except ConnectionResetError:
            print(f"电机客户端 {client_address} 强制断开连接")
        finally:
            with self.lock:
                if client_socket in self.clients:
                    self.clients.remove(client_socket)
            client_socket.close()

    def _send_to_all_clients(self):
        while True:
            command = self.command_queue.get()
            if command == "QUIT":
                break

            with self.lock:
                for client in self.clients[:]:
                    try:
                        client.send(command.encode('utf-8'))
                    except Exception as e:
                        print(f"发送电机命令失败：{e}")
                        if client in self.clients:
                            self.clients.remove(client)

    def send_command(self, command):
        self.command_queue.put(command)

    def stop(self):
        self.command_queue.put("QUIT")


class ThroughputDetector:
    """吞吐检测器 - 基于中线交点法"""

    def __init__(self, model_path='runs/detect/train3/weights/best.pt'):
        """
        初始化吞吐检测器（中线交点法）
        """
        print(f"🔍 加载训练好的模型: {model_path}")

        if not os.path.exists(model_path):
            raise FileNotFoundError(f"❌ 找不到模型文件: {model_path}")

        try:
            self.model = YOLO(model_path)
            print("✅ 模型加载成功")
        except Exception as e:
            raise RuntimeError(f"❌ 模型加载失败: {str(e)}")

        # 配置参数
        self.confidence_threshold = 0.1

        # 吞吐检测状态
        self.brightness_values = []  # 存储所有亮度值
        self.frame_numbers = []  # 存储对应的帧号
        self.crossing_points = []  # 存储中线交点位置
        self.current_throughput = 0  # 当前吞吐数量
        self.crossing_events = []  # 存储交点事件
        self.midline = 0.5  # 初始中线值

        print("✅ 中线交点法吞吐器初始化完成")

    def get_center_region_brightness(self, frame, center_point, radius=20):
        """
        获取中心区域的亮度值
        """
        x, y = int(center_point[0]), int(center_point[1])
        h, w = frame.shape[:2]

        # 确保区域在图像范围内
        x1 = max(0, x - radius)
        y1 = max(0, y - radius)
        x2 = min(w, x + radius)
        y2 = min(h, y + radius)

        if x1 >= x2 or y1 >= y2:
            return 0

        # 提取区域并计算亮度
        region = frame[y1:y2, x1:x2]
        if len(region) == 0:
            return 0

        if len(region.shape) == 3:
            region = cv2.cvtColor(region, cv2.COLOR_BGR2GRAY)

        brightness = np.mean(region) / 255.0
        return brightness

    def detect_center_point(self, frame):
        """
        检测帧中的中心点
        """
        try:
            results = self.model(frame, imgsz=640, conf=self.confidence_threshold, verbose=False)

            for result in results:
                if hasattr(result, 'boxes') and result.boxes is not None and len(result.boxes) > 0:
                    boxes = result.boxes.xyxy.cpu().numpy()
                    confidences = result.boxes.conf.cpu().numpy()

                    # 取置信度最高的检测框
                    if len(confidences) > 0:
                        max_conf_idx = np.argmax(confidences)
                        box = boxes[max_conf_idx]
                        conf = confidences[max_conf_idx]

                        # 计算中心点
                        x_center = (box[0] + box[2]) / 2
                        y_center = (box[1] + box[3]) / 2

                        return (x_center, y_center), conf

            return None, 0

        except Exception as e:
            print(f"⚠️ 中心点检测错误: {str(e)}")
            return None, 0

    def detect_midline_crossings(self):
        """
        检测亮度曲线与中线的交点
        返回：交点数量（吞吐次数 = 有效交点对数）
        过滤单个交点：只计算成对的交点（一个上升沿和一个下降沿）
        """
        if len(self.brightness_values) < 2:
            return 0, []

        # 计算中线（亮度平均值）
        self.midline = np.mean(self.brightness_values)

        # 检测所有可能的交点
        all_crossings = []
        for i in range(1, len(self.brightness_values)):
            prev_val = self.brightness_values[i - 1]
            curr_val = self.brightness_values[i]

            # 检查是否穿过中线
            if (prev_val < self.midline and curr_val >= self.midline) or \
                    (prev_val > self.midline and curr_val <= self.midline):
                all_crossings.append(i)

        # 过滤单个交点 - 只保留成对的交点
        valid_crossings = []
        i = 0
        while i < len(all_crossings) - 1:
            # 检查当前交点和下一个交点是否构成一对（一个上升一个下降）
            current_idx = all_crossings[i]
            next_idx = all_crossings[i + 1]

            # 获取交点的亮度值
            current_val = self.brightness_values[current_idx]
            next_val = self.brightness_values[next_idx]

            # 检查是否构成一对（一个高于中线，一个低于中线）
            if (current_val >= self.midline and next_val <= self.midline) or \
                    (current_val <= self.midline and next_val >= self.midline):
                valid_crossings.append(current_idx)
                valid_crossings.append(next_idx)
                i += 2  # 跳过下一对
            else:
                i += 1  # 移动到下一个点

        # 吞吐次数 = 有效交点对数
        throughput_count = len(valid_crossings) // 2

        return throughput_count, valid_crossings

    def update_throughput_count(self, frame_idx, brightness):
        """
        实时更新吞吐数量
        """
        if len(self.brightness_values) < 2:
            return self.current_throughput

        # 更新中线值
        if len(self.brightness_values) > 10:  # 至少有10个样本后再计算中线
            self.midline = np.mean(self.brightness_values)

        # 检查是否穿过中线
        prev_val = self.brightness_values[-2] if len(self.brightness_values) >= 2 else brightness
        curr_val = brightness

        # 检查是否穿过中线
        crossed = False
        direction = None

        if prev_val < self.midline and curr_val >= self.midline:
            crossed = True
            direction = "up"
        elif prev_val > self.midline and curr_val <= self.midline:
            crossed = True
            direction = "down"

        if crossed:
            # 记录交点事件
            self.crossing_events.append({
                'frame': frame_idx,
                'brightness': curr_val,
                'direction': direction,
                'time': time.time()
            })

            # 检查是否形成有效对
            if len(self.crossing_events) >= 2:
                last_event = self.crossing_events[-1]
                prev_event = self.crossing_events[-2]

                # 如果方向相反且时间间隔合理（避免误检）
                if (last_event['direction'] != prev_event['direction'] and
                        last_event['frame'] - prev_event['frame'] < 100):  # 最大100帧间隔
                    self.current_throughput += 1
                    # 移除已配对的事件
                    self.crossing_events = self.crossing_events[:-2]

        return self.current_throughput

    def draw_detection_info(self, frame, center_point, brightness, throughput_count):
        """
        在帧上绘制检测信息和吞吐数量
        """
        display_frame = frame.copy()

        # 绘制中心点
        if center_point:
            x, y = int(center_point[0]), int(center_point[1])
            cv2.circle(display_frame, (x, y), 10, (0, 255, 0), 2)
            cv2.circle(display_frame, (x, y), 20, (0, 255, 0), 1)

        # 绘制吞吐数量
        cv2.putText(display_frame, f"Throughput: {throughput_count}", (20, 50),
                    cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)

        # 绘制亮度值
        cv2.putText(display_frame, f"Brightness: {brightness:.3f}", (20, 90),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 255), 2)

        # 绘制中线值
        cv2.putText(display_frame, f"Midline: {self.midline:.3f}", (20, 130),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 0), 2)

        # 绘制帧信息
        cv2.putText(display_frame, f"Frames: {len(self.brightness_values)}", (20, 170),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 0, 255), 2)

        return display_frame

    def reset(self):
        """重置检测器状态"""
        self.brightness_values = []
        self.frame_numbers = []
        self.crossing_points = []
        self.current_throughput = 0
        self.crossing_events = []
        self.midline = 0.5


class VideoProcessor:
    def __init__(self, motor_controller, model_path='runs/detect/train3/weights/best.pt'):
        self.motor_controller = motor_controller
        self.throughput_detector = ThroughputDetector(model_path)
        self.paused = True  # 初始状态为暂停
        self.threshold = 50  # 吞吐阈值
        self.waiting_for_start = True  # 等待初始启动
        self.cap = None
        self.frame_count = 0

    def init_camera(self):
        """初始化iVCam摄像头"""
        print("正在尝试初始化iVCam摄像头...")

        # 尝试不同的索引（0, 1, 2, ...）
        self.cap = cv2.VideoCapture(0)  # 通常是0或1
        if not self.cap.isOpened():
            self.cap = cv2.VideoCapture(1)  # 如果0失败，尝试1

        if not self.cap.isOpened():
            print("无法打开摄像头，请检查iVCam是否已连接")
            return False

        # 设置摄像头分辨率
        self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1280)
        self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 720)
        self.cap.set(cv2.CAP_PROP_FPS, 30)

        print("iVCam摄像头初始化成功")
        return True

    def process_camera(self):
        """处理摄像头实时视频流"""
        if not self.init_camera():
            return

        ret, first_frame = self.cap.read()
        if not ret:
            print("无法从摄像头读取帧")
            return

        height, width = first_frame.shape[:2]

        # 初始状态：电机停止，视频暂停
        self.motor_controller.send_command("stop")
        self.paused = True
        self.waiting_for_start = True
        self.throughput_detector.reset()

        print("系统就绪，按 's' 开始检测，按 'q' 退出")

        while True:
            # 显示等待启动提示
            if self.waiting_for_start:
                display_frame = np.zeros((height, width, 3), dtype=np.uint8)
                cv2.putText(display_frame, "Press 's' to start motor and detection",
                            (width // 6, height // 2 - 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
                cv2.putText(display_frame, "Press 'q' to quit",
                            (width // 3, height // 2 + 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
                cv2.imshow('Real-time Throughput Detection', display_frame)

            # 处理暂停状态下的人工启动
            if self.paused and not self.waiting_for_start:
                display_frame = np.zeros((height, width, 3), dtype=np.uint8)
                cv2.putText(display_frame, "PAUSED - Press 's' to continue",
                            (width // 4, height // 2 - 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
                cv2.putText(display_frame,
                            f"Throughput Count: {self.throughput_detector.current_throughput}/{self.threshold}",
                            (width // 4, height // 2 + 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
                cv2.imshow('Real-time Throughput Detection', display_frame)

            # 处理键盘输入
            key = cv2.waitKey(1) & 0xFF
            if key == ord('q'):
                break
            elif key == ord('s'):
                if self.waiting_for_start:
                    # 初始启动
                    self.motor_controller.send_command("start")
                    self.paused = False
                    self.waiting_for_start = False
                    self.throughput_detector.reset()
                    print("开始检测，电机启动")
                elif self.paused:
                    # 暂停后恢复
                    self.throughput_detector.reset()
                    self.motor_controller.send_command("start")
                    self.paused = False
                    print("恢复检测，电机启动")

            # 正常处理帧
            if not self.paused and not self.waiting_for_start:
                ret, frame = self.cap.read()
                if not ret:
                    print("无法从摄像头获取帧")
                    break

                self.frame_count += 1

                # 检测中心点
                center_point, confidence = self.throughput_detector.detect_center_point(frame)

                brightness = 0
                if center_point:
                    # 获取中心区域亮度
                    brightness = self.throughput_detector.get_center_region_brightness(frame, center_point)
                    self.throughput_detector.brightness_values.append(brightness)
                    self.throughput_detector.frame_numbers.append(self.frame_count)

                    # 实时更新吞吐数量
                    throughput_count = self.throughput_detector.update_throughput_count(
                        self.frame_count, brightness)

                    # 检查是否达到阈值
                    if throughput_count >= self.threshold:
                        self.motor_controller.send_command("stop")
                        self.paused = True
                        print(f"达到吞吐阈值 {self.threshold}，电机停止")
                else:
                    # 如果没有检测到中心点，使用上一次的吞吐计数
                    throughput_count = self.throughput_detector.current_throughput

                # 显示检测结果
                display_frame = self.throughput_detector.draw_detection_info(
                    frame, center_point, brightness, throughput_count)

                # 添加状态信息
                status = "RUNNING" if not self.paused else "PAUSED"
                cv2.putText(display_frame, f"Status: {status}", (20, 210),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 0, 0), 2)
                cv2.putText(display_frame, f"Threshold: {self.threshold}", (20, 250),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)

                # 调整显示窗口大小
                display_height = min(800, height)
                scale_factor = display_height / height
                display_width = int(width * scale_factor)
                resized_frame = cv2.resize(display_frame, (display_width, display_height))

                cv2.imshow('Real-time Throughput Detection', resized_frame)

        self.cap.release()
        cv2.destroyAllWindows()

    def generate_report(self):
        """生成检测报告"""
        if len(self.throughput_detector.brightness_values) > 0:
            # 最终的中线交点检测（更准确）
            final_throughput_count, crossings = self.throughput_detector.detect_midline_crossings()

            print(f"\n📊 最终检测结果:")
            print(f"总帧数: {self.frame_count}")
            print(f"有效亮度数据帧数: {len(self.throughput_detector.brightness_values)}")
            print(f"中线值: {self.throughput_detector.midline:.3f}")
            print(f"实时吞吐数量: {self.throughput_detector.current_throughput}")
            print(f"最终吞吐数量: {final_throughput_count}")

            return final_throughput_count
        return 0


def main():
    # 初始化电机控制器
    motor_controller = MotorController()
    motor_controller.start_server()

    # 初始化视频处理器和吞吐检测器
    # 请确保修改为正确的模型路径
    model_path = 'runs/detect/train3/weights/best.pt'  # 修改为你的模型路径
    video_processor = VideoProcessor(motor_controller, model_path)

    try:
        # 开始处理摄像头视频流
        video_processor.process_camera()

        # 生成最终报告
        video_processor.generate_report()

    except Exception as e:
        print(f"发生错误: {e}")
    finally:
        # 确保电机停止
        motor_controller.send_command("stop")
        motor_controller.stop()
        print("程序结束")


if __name__ == "__main__":
    main()